# -*- coding: utf-8 -*-
import requests
import json
import logging
import os
import glob



# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class LLMOCR:


    def __init__(self, file_path: str):
        self.dify_url = "http://10.237.14.79/v1/workflows/run"  # 替换为实际的 Dify 接口 URL
        self.api_key = "app-Vz9gjPE171hhrDZI3NbZUOto"
        self.file_path = file_path

    def upload_file(self) -> dict:
        """上传文件并返回文件信息"""
        url = "http://10.237.14.79/v1/files/upload"
        headers = {
            "Authorization": f"Bearer {self.api_key}"
        }
        
        # 调试输出文件信息
        file_extension = os.path.splitext(self.file_path)[1].lower()
        mime_type = "image/jpeg"  # 默认 MIME 类型
        if file_extension in ['.png']:
            mime_type = 'image/png'
        elif file_extension in ['.jpg', '.jpeg']:
            mime_type = 'image/jpeg'
        elif file_extension in ['.webp']:
            mime_type = 'image/webp'
        elif file_extension in ['.gif']:
            mime_type = 'image/gif'
            
        logger.info(f"准备上传文件: {self.file_path}")
        logger.info(f"文件扩展名: {file_extension}")
        logger.info(f"MIME 类型: {mime_type}")
        
        with open(self.file_path, 'rb') as f:
            files = {
                'file': (os.path.basename(self.file_path), f, mime_type)
            }
            data = {
                'user': 'abc-123'
            }
            response = requests.post(url, headers=headers, files=files, data=data)
            
            if response.status_code in [200, 201]:
                response_data = response.json()
                logger.info(f"File uploaded successfully: {response_data}")
                return response_data  # 返回文件信息
            else:
                logger.error(f"Failed to upload file. Status code: {response.status_code}")
                logger.error(f"Response: {response.text}")
                return None

    def run_workflow(self) -> str:
        """运行工作流"""
        # 上传文件并获取文件信息
        upload_response = self.upload_file()
        if not upload_response:
            return False  # 上传失败，返回

        logger.info(f"上传响应: {upload_response}")
        # 从上传响应中提取文件信息
        upload_file_id = upload_response['id']
        filename = upload_response['name']
        size = upload_response['size']
        extension = upload_response['extension']
        mime_type = upload_response['mime_type']
        related_id = upload_response['id']  # 假设 related_id 是上传返回的 ID
        # preview_url = f"/files/{related_id}/file-preview?timestamp=1740278189&nonce=0919fbb2a088fc6cafca1d90e866f392&sign=SIBdaVKi40U_7J19NHzfMjXTN35hofKoRNGRRuZ9uLI="
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        # 准备 JSON 数据
        data = {
            "inputs": {
                "input_text": {
                    "dify_model_identity": "__dify__file__",
                    "id": None,
                    "tenant_id": "8e1efdf4-ca44-4ab8-8ba0-717344c51f5b",
                    "type": "image",
                    "transfer_method": "local_file",
                    "remote_url": "",
                    "related_id": related_id,
                    "filename": filename,
                    "extension": f".{extension}",  # 添加点号
                    "mime_type": mime_type,
                    "size": size,
                    "upload_file_id": upload_file_id,
                    # "url": preview_url
                }
            },
            "response_mode": "blocking",
            "user": "abc-123",
            "sys.files": [],
            "sys.user_id": "ff7f33d3-26e8-44fb-8b59-97e2fcc4c3fd",
            "sys.app_id": "5e21d896-a56d-4cff-80cf-3d107aef8974",
            "sys.workflow_id": "43edcf5f-3c69-479a-8382-1192ad9962ae",
            "sys.workflow_run_id": "8329cf44-5a14-4fa4-a8e8-08c6a702c680"
        }
        print(data)
        
        # 发送 POST 请求到工作流接口
        response = requests.post(self.dify_url, headers=headers, json=data)
        # 处理响应...
        if response.status_code == 200:
            response_data = response.json()
            logger.info(f"[LLM Response]")
            logger.info(f"Raw response: {json.dumps(response_data, ensure_ascii=False, indent=2)}")
            return response_data
        else:
            logger.error(f"Failed to run workflow. Status code: {response.status_code}")
            logger.error(f"Response: {response.text}")
            return False

def process_all_images():
    """处理所有merged_result开头的图片并保存结果"""
    # 获取pic目录下所有merged_result开头的图片
    pic_dir = "pic"
    image_pattern = os.path.join(pic_dir, "merged_result*.jpg")
    image_files = sorted(glob.glob(image_pattern))
    
    # 创建结果文件
    result_file = "ocr_results.txt"
    
    print(f"找到 {len(image_files)} 个图片文件")
    
    with open(result_file, "w", encoding="utf-8") as f:
        for idx, image_path in enumerate(image_files, 1):
            print(f"\n处理第 {idx}/{len(image_files)} 个图片: {image_path}")
            
            # 创建OCR实例并执行
            ocr = LLMOCR(image_path)
            result = ocr.run_workflow()
            
            if result and 'data' in result and 'outputs' in result['data']:
                # 写入图片路径作为标题
                f.write(f"\n{'='*50}\n")
                f.write(f"图片: {os.path.basename(image_path)}\n")
                f.write(f"{'='*50}\n\n")
                
                # 写入识别的文本
                text = result['data']['outputs'].get('text', '')
                f.write(text)
                f.write("\n")
                
                print(f"成功处理图片: {image_path}")
            else:
                print(f"处理图片失败: {image_path}")
                # 写入错误信息
                f.write(f"\n{'='*50}\n")
                f.write(f"图片: {os.path.basename(image_path)} (处理失败)\n")
                f.write(f"{'='*50}\n\n")
        #删除image_files中的图片
        for image_path in image_files:
            os.remove(image_path)
    
    print(f"\n所有结果已保存到: {result_file}")

if __name__ == "__main__":
    # 添加命令行参数处理
    import sys
    
    if len(sys.argv) > 1:
        # 如果提供了图片路径参数，处理单个图片
        image_path = sys.argv[1]
        print(f"\n开始处理图片: {image_path}")
        
        # 创建 OCR 实例并执行
        ocr = LLMOCR(image_path)
        result = ocr.run_workflow()
        
        if result:
            print("\n识别结果:")
            print("-" * 50)
            print(result)
            print("-" * 50)
        else:
            print("\n识别失败，请检查日志获取详细信息")
    else:
        # 如果没有提供参数，处理所有图片
        process_all_images()





